#!/bin/bash

============ 全局变量定义 ============
SSHD_CONFIG="/etc/ssh/sshd_config"
BACKUP_DIR="/etc/ssh/backup"
TIMESTAMP=$(date +"%Y%m%d%H%M%S")
LOG_FILE="/var/log/ssh_secure.log"

============ 日志记录函数 ============
log() {
    echo "$(date '+%Y-%m-%d %H:%M:%S') [INFO] $1" | tee -a "$LOG_FILE"
}

error_log() {
    echo "$(date '+%Y-%m-%d %H:%M:%S') [ERROR] $1" | tee -a "$LOG_FILE"
}

============ 备份配置文件 ============
backup_sshd_config() {
    if [ ! -d "$BACKUP_DIR" ]; then
        mkdir -p "$BACKUP_DIR"
    fi
    cp "$SSHD_CONFIG" "$BACKUP_DIR/sshd_config.bak.$TIMESTAMP"
    log "备份 SSH 配置文件到: $BACKUP_DIR/sshd_config.bak.$TIMESTAMP"
}

============ 修改 SSH 配置 ============
configure_sshd() {
    log "开始修改 SSH 配置..."

    # 备份原有配置
    backup_sshd_config

    # 替换或添加关键配置项
    sed -i 's/#PermitRootLogin yes/PermitRootLogin no/' "$SSHD_CONFIG"
    echo "PermitRootLogin no" >> "$SSHD_CONFIG"

    echo "PasswordAuthentication no" >> "$SSHD_CONFIG"
    echo "PubkeyAuthentication yes" >> "$SSHD_CONFIG"
    echo "PermitEmptyPasswords no" >> "$SSHD_CONFIG"
    echo "AllowTcpForwarding no" >> "$SSHD_CONFIG"
    echo "GatewayPorts no" >> "$SSHD_CONFIG"
    echo "X11Forwarding no" >> "$SSHD_CONFIG"
    echo "MaxAuthTries 3" >> "$SSHD_CONFIG"
    echo "LogLevel VERBOSE" >> "$SSHD_CONFIG"
    echo "ClientAliveInterval 300" >> "$SSHD_CONFIG"
    echo "ClientAliveCountMax 0" >> "$SSHD_CONFIG"

    log "SSH 配置修改完成。"
}

============ 创建登录白名单用户组 ============
create_ssh_users_group() {
    read -p "请输入允许 SSH 登录的用户名列表（空格分隔）：" USERS
    for USER in $USERS; do
        if id "$USER" &>/dev/null; then
            echo "AllowUsers $USER" >> "$SSHD_CONFIG"
        else
            error_log "用户 $USER 不存在，跳过。"
        fi
    done
    log "已设置 AllowUsers 白名单。"
}

============ 配置 PAM 登录限制 ============
configure_pam_limits() {
    log "配置 PAM 登录限制..."
    echo "auth required pam_access.so" >> /etc/pam.d/sshd
    echo "auth required pam_tally2.so deny=5 unlock_time=600" >> /etc/pam.d/sshd
    log "PAM 登录限制配置完成。"
}

============ 配置防火墙规则（可选） ============
configure_firewall() {
    read -p "是否启用防火墙限制 SSH 源地址？(y/n): " SET_FIREWALL
    if [[ "$SET_FIREWALL" == "y" ]]; then
        read -p "请输入允许访问 SSH 的 IP 地址范围（如 192.168.1.0/24 或多个IP空格分隔）：" IPS
        for IP in $IPS; do
            iptables -A INPUT -p tcp --dport 22 -s "$IP" -j ACCEPT
        done
        iptables -A INPUT -p tcp --dport 22 -j DROP
        service iptables save || true
        log "防火墙规则已限制 SSH 源地址。"
    fi
}

============ 重启 SSHD 服务 ============
restart_sshd() {
    log "重启 SSH 服务..."
    systemctl restart sshd
    if systemctl is-active --quiet sshd; then
        log "SSH 服务重启成功。"
    else
        error_log "SSH 服务重启失败，请手动检查。"
        exit 1
    fi
}

============ 主函数 ============
main() {
    log "开始执行 SSH 安全加固脚本..."

    configure_sshd
    create_ssh_users_group
    configure_pam_limits
    configure_firewall
    restart_sshd

    log "SSH 安全加固已完成！请确保保留当前终端连接，避免被锁在外面。"
}

main